#include "armordetector.hpp"

namespace AD{
    using namespace cv;
    void Armor_Detector::detect(){
        If_Detect=true;
        img=src.clone();
        Threshold();
        find_light_bar();
        find_armor();
    }
    void Armor_Detector::Threshold(){
        GaussianBlur(img,binary_img,Size(5,5),5,5);
        cvtColor(binary_img,binary_img,COLOR_BGR2HSV);
        for(register int i=0; i<binary_img .rows; i++)
            for(register int j=0; j<binary_img.cols; j++){
                Vec3b &t=binary_img.at<Vec3b>(i,j);
                if(t[0]>=80 && t[0]<=110 && t[1]>=40 && t[2]>=100)
                    t[0]=t[1]=t[2]=255;
                else
                    t[0]=t[1]=t[2]=0;
            }
        cvtColor(binary_img,binary_img,COLOR_BGR2GRAY);
        Mat kernel1=getStructuringElement(MORPH_CROSS,Size(14,20));
        dilate(binary_img,binary_img,kernel1);
        erode (binary_img,binary_img,kernel1);
        Mat kernel2=getStructuringElement(MORPH_CROSS,Size(8 ,10));
        erode (binary_img,binary_img,kernel2);
        dilate(binary_img,binary_img,kernel2);
        findContours(binary_img,contours,RETR_EXTERNAL,CHAIN_APPROX_SIMPLE);
        imshow("h",binary_img);
    }

    double Armor_Detector::get_angle(Point2f a,Point2f b){
        double dy=a.x-b.x,dx=a.y-b.y;
        double angle=atan(dy/dx);
        if(dx<0 && dy>0)angle+=PI;
        if(dx<0 && dy<0)angle+=PI;
        if(dx>0 && dy<0)angle+=PI*2;
        if(angle>PI)angle-=PI;
        return 180 * angle / PI;
    }

    void Armor_Detector::find_light_bar(){
        for(int i=0;i<contours.size();i++){
            if(contourArea(contours[i])<20)continue;
            RotatedRect rect=minAreaRect(contours[i]);
            if(contourArea(contours[i])/(rect.size.width*rect.size.height)>0.7){
                if(max(rect.size.width,rect.size.height)/min(rect.size.width,rect.size.height)<3)continue;
                draw_light_bar(rect);
                light_bar.push_back(rect);
            }
        }
    }

    void Armor_Detector::find_armor(){
        for(int i=0;i<light_bar.size();i++){
            for(int j=i+1;j<light_bar.size();j++){
                RotatedRect a=light_bar[i],b=light_bar[j];
                double k1=a.angle,k2=b.angle;
                double k3=get_angle(a.center,b.center);
                if(k3>135 || k3<45)continue;
                double t1=(k1+k3)/90,t2=(k2+k3)/90;
                if((t1-floor(t1)<0.1 || ceil(t1)-t1<0.1) && (t2-floor(t2)<0.1 || ceil(t2)-t2<0.1))
                    draw_armor(a,b);
            }
        }
    }
    void Armor_Detector::draw_light_bar(RotatedRect &rect){
        Point2f rectPoints[4];
        rect.points(rectPoints);
        for(int i=0; i<4; i++){
            line(img,rectPoints[i],rectPoints[(i+1)%4],Scalar(0,255,0),2);
        }
    }

    void Armor_Detector::draw_armor(RotatedRect &a, RotatedRect &b){
        Point2f rectPoints[8];
        std::vector<Point2f> armorPoints(4);
        a.points(rectPoints  );
        b.points(rectPoints+4);
        double val[4];
        for(int i=0;i<8;i++){
            Point2f now = rectPoints[i];
            if(i==0 || -now.x-now.y>val[0]) val[0]=-now.x-now.y, armorPoints[0]=now;
            if(i==0 ||  now.x-now.y>val[1]) val[1]= now.x-now.y, armorPoints[1]=now;
            if(i==0 ||  now.x+now.y>val[2]) val[2]= now.x+now.y, armorPoints[2]=now;
            if(i==0 || -now.x+now.y>val[3]) val[3]=-now.x+now.y, armorPoints[3]=now;
        }
        for(int i=0;i<4;i++){
            line(img,armorPoints[i],armorPoints[(i+1)%4],Scalar(0,0,255),3);
        }
        Point2f target=(a.center+b.center)/2;
        circle(img,target,5,Scalar(0,0,255),-1);
        armor.push_back(armorPoints);
    }

    void Armor_Detector::calc_dis(){
        if(!If_Detect)detect();
        If_CalcDis=true;
        for(int i=0;i<armor.size();i++){
            Point2f rectPoints[4];
            std::vector<Point2f> armorInImage(4);
            armorInImage=armor[i];
            Mat rvec, tvec, rotateMat;
            solvePnP(armorSize,armorInImage,cameraMatrix,distCoeffs,rvec,tvec);
            double x= tvec.at<double>(0,0);
            double y= tvec.at<double>(1,0);
            double z= tvec.at<double>(2,0);
            armorInCamera.push_back(Point3f(x,y,z));
        }
    }

    void Armor_Detector::draw_dis(){
        for(int i=0;i<armor.size();i++){
            If_DrawDis=true;
            Point2f target = (armor[i][0]+armor[i][2])/2;
            Point3f coordinate = armorInCamera[i];
            double depth= coordinate.z;
            double dis= sqrt(coordinate.x*coordinate.x + coordinate.y*coordinate.y + coordinate.z*coordinate.z);
            putText(img,"dis="  +std::to_string(  dis)+"cm",target+Point2f(35,0 ),FONT_HERSHEY_SIMPLEX,0.7,Scalar(0,0,0),3,LINE_AA);
            putText(img,"depth="+std::to_string(depth)+"cm",target+Point2f(0,-25),FONT_HERSHEY_SIMPLEX,0.7,Scalar(0,0,0),3,LINE_AA);
            putText(img,"dis="  +std::to_string(  dis)+"cm",target+Point2f(35,0 ),FONT_HERSHEY_SIMPLEX,0.7,Scalar(255,255,255),1,LINE_AA);
            putText(img,"depth="+std::to_string(depth)+"cm",target+Point2f(0,-25),FONT_HERSHEY_SIMPLEX,0.7,Scalar(255,255,255),1,LINE_AA);
        }
    }
}